# TODO(b/157939296): This file is required for running using XCloud.
absl-py
clu
flax

jax==0.3.4
--find-links https://storage.googleapis.com/jax-releases/jax_releases.html
jaxlib==0.3.2+cuda11.cudnn82  # Make sure CUDA version matches the base image.

ml_collections
numpy
tensorflow-cpu
tensorflow-datasets
